import matplotlib.pyplot as plt
import pandas as pd

if __name__ == "__main__":
    df = pd.read_csv("results_others.csv").set_index('index')

    # data copied from other scripts
    df.loc['markovstab', 'vi'] = 0.1863
    df.loc['markovstab', 'n_comms'] = 245
    df.loc['geocluster', 'vi'] = 0.0852
    df.loc['geocluster', 'n_comms'] = 154

    df.loc['wavelet', 'vi'] = 0.24
    df.loc['wavelet', 'n_comms'] = 49
    df.loc['wavelet_best', 'vi'] = 0.147
    df.loc['wavelet_best', 'n_comms'] = 94

    df = df.sort_values(by='vi')
    df = df[:-2]
    plt.figure(figsize=(6, 3))
    df['vi'].plot.bar(ax=plt.gca())
    plt.savefig('com_clustering_vi.pdf', bbox_inches='tight')

    plt.figure(figsize=(6, 3))
    df['n_comms'].plot.bar(ax=plt.gca())
    plt.axhline(302, label='number of nodes', c='k')
    plt.axhline(118, label='number of groundtruth clusters', c='r')
    plt.legend()
    plt.savefig('com_clustering_n_comms.pdf', bbox_inches='tight')
